import json
import random
import re
import string
import time
import traceback
import uuid
from subprocess import TimeoutExpired
from typing import List

import requests
import yaml
from minisweagent.environments.docker import DockerEnvironment, DockerEnvironmentConfig
from openai import OpenAI
from swebench.harness.run_evaluation import run_instance
from swebench.harness.test_spec.test_spec import make_test_spec, TestSpec

from minisweagent.run.extra.swebench import (
    DATASET_MAPPING,
    get_sb_environment,
)
import docker
from minisweagent.run.extra.swebench import get_swebench_docker_image_name


class SWEBenchShellEnv(object):
    """Gym environment for HTML mode of WebShop environment"""

    def __init__(self,**kwargs):
        
        config_path = "/usr/local/python3.11.13/lib/python3.11/site-packages/minisweagent/config/extra/swebench.yaml"
        # config_path = "/home/tianye/code/mini-swe-agent/src/minisweagent/config/extra/swebench.yaml"
        self.config = yaml.safe_load("\n".join(open(config_path, "r+").readlines()))
        self.env = None
        # self.reset()

    def step(self, action, valid):
        """
        Takes an action, updates WebShop environment, and returns (observation, reward, done, info)

        Arguments:
        action (`str`): An action should be of the following structure:
          - search[keywords]
          - click[value]
        If action not valid, perform nothing.
        """
        # observation = " "
        # reward = 0.0
        # done = False
        # info = {}
        # info["available_actions"] = self.obtain_api_list(self.index)
        # self._step += 1
        # ref" current step {self._step} , max step {self.max_step}")
        # done = self._step >= self.max_step
        if valid:
            try:
                if isinstance(action, str):
                    res = self.env.execute(action, None)
                else:
                    res = self.env.execute(action.get("command", ""), action.get("cwd",None))
            except TimeoutExpired:
                return "command exec time out", 0 , False, {}

            info = {"returncode":res["returncode"]}

            return res["output"], 1 if res["returncode"] == 0 else 0, False, info
        else:
            return "action is not valid.", -1, True, {}

    @property
    def action_space(self):
        # Recommended to use `get_available_actions` instead
        return NotImplementedError

    @property
    def observation_space(self):
        return NotImplementedError

    def reset(self, instance):
        #add observation and info as self
        """Create a new session and reset environment variables"""
        # self.observation = self.obtain_api_list(index)
        if self.env is not None:
            self.env.cleanup()
        iid = instance["instance_id"]
        id_docker_compatible = iid.replace("__","_1776_")
        instance["image_name"] = instance["image_name"] if "image_name" in instance else f"swebench/sweb.eval.x86_64.{id_docker_compatible}:latest".lower()
        instance["image_name"] = instance["image_name"].replace("x86_64","arm64")
        image_name = get_swebench_docker_image_name(instance)
        self.env = get_sb_environment(self.config, instance)
        # instance["image_name"] = "docker.xuanyuan.run/"+instance["image_name"] if "image_name" in instance else None

        # self.env =get_sb_environment() DockerEnvironment(**(self.config.get("environment", {}) | {"image": image_name}))
        self.instance = instance

        return instance["problem_statement"], {}

    def render(self, mode='human'):
        # TODO: Render observation in terminal or WebShop website
        return NotImplementedError

    def close(self):
        # TODO: When DB used instead of JSONs, tear down DB here
        if self.env is not None:
            self.env.cleanup()
